
def main4():
    import cv2
    import numpy as np
    import os
    import random
    def load_images(folder, is_mask=False):
        '''is_mask=True 时，会将掩码图像转换为单通道（灰度）模式，以便与 mosaic_msk 匹配。这可以避免形状不匹配的问题。'''
        files = os.listdir(folder)
        images = []
        for file in files:
            img_path = os.path.join(folder, file)
            img = cv2.imread(img_path)
            if img is not None:
                if is_mask:
                    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # 将掩码转换为单通道
                images.append(img)
            else:
                print(f"Error loading image: {file}")
        return images

    # 随机裁剪函数
    def random_crop(image, mask, crop_size=(100, 100)):
        # print(image.shape)
        h, w = image.shape[:2]
        crop_h, crop_w = crop_size

        # 确保裁剪区域不会超出边界
        top = random.randint(0, h - crop_h)
        left = random.randint(0, w - crop_w)

        crop_img = image[top:top + crop_h, left:left + crop_w]
        crop_mask = mask[top:top + crop_h, left:left + crop_w]

        return crop_img, crop_mask

    # Mosaic数据增强函数
    def mosaic_augmentation(images, masks):
        # 随机选择四张图像和对应掩码
        idxs = random.sample(range(len(images)), 4)
        imgs = [images[i] for i in idxs]
        msks = [masks[i] for i in idxs]

        # 创建200x200的空白Mosaic图像和掩码
        mosaic_img = np.zeros((200, 200, 3), dtype=np.uint8)
        mosaic_msk = np.zeros((200, 200), dtype=np.uint8)

        # 将四张裁剪后的图像和掩码合并到200x200的图像中
        mosaic_img[0:100, 0:100] = imgs[0]
        mosaic_img[0:100, 100:200] = imgs[1]
        mosaic_img[100:200, 0:100] = imgs[2]
        mosaic_img[100:200, 100:200] = imgs[3]

        mosaic_msk[0:100, 0:100] = msks[0]
        mosaic_msk[0:100, 100:200] = msks[1]
        mosaic_msk[100:200, 0:100] = msks[2]
        mosaic_msk[100:200, 100:200] = msks[3]

        return mosaic_img, mosaic_msk

    # 文件夹路径
    image_folder_test = r'./data/neu_seg_competition/images/test'
    mask_folder_test = r'./data/neu_seg_competition/annotations/test'
    image_folder_training = r'./data/neu_seg_competition/images/training'
    mask_folder_training = r'./data/neu_seg_competition/annotations/training'
    show = 0  # 展示图片增强结果？ 1：True  0:False
    # 加载图像和掩码
    images = load_images(image_folder_test)
    masks = load_images(mask_folder_test, is_mask=True)  # 指定is_mask=True将掩码转换为单通道
    images_1 = load_images(image_folder_training)
    masks_1 = load_images(mask_folder_training, is_mask=True)

    # 合并训练集和测试集数据
    images += images_1
    masks += masks_1

    # 随机裁剪图像和掩码到100x100大小
    cropped_images = []
    cropped_masks = []

    for img, msk in zip(images, masks):
        crop_img, crop_msk = random_crop(img, msk, crop_size=(100, 100))
        cropped_images.append(crop_img)
        cropped_masks.append(crop_msk)

    if show:
        # 执行Mosaic数据增强
        mosaic_img, mosaic_msk = mosaic_augmentation(cropped_images, cropped_masks)

        # 显示Mosaic增强结果
        cv2.imshow("Mosaic Image", mosaic_img)
        cv2.imshow("Mosaic Mask", mosaic_msk*225)  # 放大显示掩码
        cv2.waitKey(0)
        cv2.destroyAllWindows()


    import os
    import shutil
    if not os.path.exists('./data/neu_seg_competition_nation_self/images/training/'):
        os.makedirs('./data/neu_seg_competition_nation_self/images/training/')

    if not os.path.exists('./data/neu_seg_competition_nation_self/annotations/training/'):
        os.makedirs('./data/neu_seg_competition_nation_self/annotations/training/')


    # 离线增强保存（可选择启用）
    for i in range(500):
        mosaic_img, mosaic_msk = mosaic_augmentation(cropped_images, cropped_masks)
        cv2.imwrite(f'./data/neu_seg_competition_nation_self/images/training/4_mosaic_{i}.jpg', mosaic_img)
        cv2.imwrite(f'./data/neu_seg_competition_nation_self/annotations/training/4_mosaic_{i}.png', mosaic_msk)
        # print(f'成功 {i+1}')
